package com.timeriver.cases.regression

import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.evaluation.{RegressionEvaluator}
import org.apache.spark.ml.feature.{StandardScaler, VectorAssembler}
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.regression.DecisionTreeRegressor
import org.apache.spark.ml.tuning.{CrossValidator, CrossValidatorModel, ParamGridBuilder}
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.types.DoubleType
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * 波士顿房屋价格数据回归预测
  *   rmse===>4.305998642096831
  *   mse===>18.541624305739752
  *   r2===>0.8088814876509002
  *   mae===>2.8754250217727853
  */
object DecisionTreeRegression {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .master("local[6]")
      .appName("决策树回归算法")
      .getOrCreate()

    val rawData: DataFrame = session.read
      .format("csv")
      .option("header", "false")
      .option("delimiter", ",")
      .load("D:\\workspace\\gitee_space\\spark-ml-machine-learning\\data\\housing.csv")

    val data: DataFrame = rawData.toDF("CRIM", "ZN", "INDUS", "CHAS", "NOX", "RM", "AGE"
      , "DIS", "RAD", "TAX", "PTRATIO", "B", "LSTAT", "label")

    val trainData: DataFrame = data.withColumn("CRIM", col("CRIM").cast(DoubleType))
      .withColumn("ZN", col("ZN").cast(DoubleType))
      .withColumn("INDUS", col("INDUS").cast(DoubleType))
      .withColumn("CHAS", col("CHAS").cast(DoubleType))
      .withColumn("NOX", col("NOX").cast(DoubleType))
      .withColumn("RM", col("RM").cast(DoubleType))
      .withColumn("AGE", col("AGE").cast(DoubleType))
      .withColumn("DIS", col("DIS").cast(DoubleType))
      .withColumn("RAD", col("RAD").cast(DoubleType))
      .withColumn("TAX", col("TAX").cast(DoubleType))
      .withColumn("PTRATIO", col("PTRATIO").cast(DoubleType))
      .withColumn("B", col("B").cast(DoubleType))
      .withColumn("LSTAT", col("LSTAT").cast(DoubleType))
      .withColumn("label", col("label").cast(DoubleType))

    /** 数据基本情况统计 */
    trainData.summary().show(false)

    val assembler: DataFrame = new VectorAssembler()
      .setInputCols(Array("CRIM", "ZN", "INDUS", "CHAS", "NOX", "RM", "AGE", "DIS", "RAD", "TAX", "PTRATIO", "B", "LSTAT"))
      .setOutputCol("features")
      .transform(trainData)

    /** 切分数据集合 */
    val Array(train, test) = assembler.randomSplit(Array(0.8, 0.2))

    val scaler: StandardScaler = new StandardScaler()
      .setInputCol("features")
      .setOutputCol("scaledFeatures")

    val treeRegressor = new DecisionTreeRegressor()

    val pipeline: Pipeline = new Pipeline().setStages(Array(scaler, treeRegressor))

    /** 构建参数网格搜索：参数调优 */
    val paramMaps: Array[ParamMap] = new ParamGridBuilder()
      .addGrid(treeRegressor.maxDepth, Array(10, 20, 30))
      .addGrid(treeRegressor.minInfoGain, Array(0.0, 0.01, 0.1, 0.3))
      .addGrid(scaler.withStd, Array(true, false))
      .addGrid(scaler.withMean, Array(true, false))
      .build()

    val validator: CrossValidator = new CrossValidator()
      .setEstimator(pipeline)
      .setEvaluator(new RegressionEvaluator)
      .setEstimatorParamMaps(paramMaps)
      /** 在生产实践中使用3或以上 */
      .setNumFolds(2)
      .setParallelism(2)

    /** 模型训练 */
    val model: CrossValidatorModel = validator.fit(train)

    /** 模型预测 */
    val res: DataFrame = model.transform(test)
    res.show(5, false)

    /** 分析报告 */
    val evaluator: RegressionEvaluator = new RegressionEvaluator()
      .setPredictionCol("prediction")
      .setLabelCol("label")
      .setMetricName("rmse")
    val rmse: Double = evaluator.evaluate(res)
    println(s"rmse===>$rmse")

    evaluator.setMetricName("mse")
    val mse: Double = evaluator.evaluate(res)
    println(s"mse===>$mse")

    evaluator.setMetricName("r2")
    val r2: Double = evaluator.evaluate(res)
    println(s"r2===>$r2")

    evaluator.setMetricName("mae")
    val mae: Double = evaluator.evaluate(res)
    println(s"mae===>$mae")

    session.stop()
  }
}
